package com.leon.datalink.core.storage;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;

/**
 * H2数据库存储实现，替换LMDB
 * 
 * @author DataLink Team
 * @version 1.0.0
 */
public class H2Storage<T> implements AutoCloseable {
    
    private static final Logger logger = LoggerFactory.getLogger(H2Storage.class);
    
    private final String dbPath;
    private final String tableName;
    private Connection connection;
    private final Class<T> valueType;
    
    // 连接池缓存
    private static final Map<String, Connection> connectionPool = new ConcurrentHashMap<>();
    
    public H2Storage(Class<T> valueType) throws SQLException {
        this(valueType, "datalink", "storage");
    }
    
    public H2Storage(Class<T> valueType, String dbName, String tableName) throws SQLException {
        this.valueType = valueType;
        this.dbPath = "jdbc:h2:./data/" + dbName + ";AUTO_SERVER=TRUE;DB_CLOSE_DELAY=-1";
        this.tableName = tableName;
        
        initializeDatabase();
    }
    
    private void initializeDatabase() throws SQLException {
        // 获取或创建连接
        this.connection = connectionPool.computeIfAbsent(dbPath, path -> {
            try {
                Connection conn = DriverManager.getConnection(path, "sa", "");
                conn.setAutoCommit(true);
                return conn;
            } catch (SQLException e) {
                logger.error("Failed to create H2 connection: {}", path, e);
                throw new RuntimeException(e);
            }
        });
        
        // 创建表
        createTableIfNotExists();
        
        logger.info("H2Storage initialized: dbPath={}, tableName={}", dbPath, tableName);
    }
    
    private void createTableIfNotExists() throws SQLException {
        String createTableSQL = String.format(
            "CREATE TABLE IF NOT EXISTS %s (" +
            "key VARCHAR(255) PRIMARY KEY, " +
            "value BLOB, " +
            "created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, " +
            "updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP" +
            ")", tableName
        );
        
        try (Statement stmt = connection.createStatement()) {
            stmt.execute(createTableSQL);
        }
        
        logger.debug("Table {} created or already exists", tableName);
    }
    
    /**
     * 存储键值对
     */
    public void put(String key, T value) throws Exception {
        byte[] serializedValue = serialize(value);
        
        String sql = String.format(
            "MERGE INTO %s (key, value, updated_at) VALUES (?, ?, CURRENT_TIMESTAMP)", 
            tableName
        );
        
        try (PreparedStatement pstmt = connection.prepareStatement(sql)) {
            pstmt.setString(1, key);
            pstmt.setBytes(2, serializedValue);
            pstmt.executeUpdate();
        }
        
        logger.debug("Stored key: {}", key);
    }
    
    /**
     * 获取值
     */
    public T get(String key) throws Exception {
        String sql = String.format("SELECT value FROM %s WHERE key = ?", tableName);
        
        try (PreparedStatement pstmt = connection.prepareStatement(sql)) {
            pstmt.setString(1, key);
            
            try (ResultSet rs = pstmt.executeQuery()) {
                if (rs.next()) {
                    byte[] serializedValue = rs.getBytes("value");
                    return deserialize(serializedValue);
                }
            }
        }
        
        return null;
    }
    
    /**
     * 删除键
     */
    public boolean delete(String key) throws SQLException {
        String sql = String.format("DELETE FROM %s WHERE key = ?", tableName);
        
        try (PreparedStatement pstmt = connection.prepareStatement(sql)) {
            pstmt.setString(1, key);
            int rowsAffected = pstmt.executeUpdate();
            
            logger.debug("Deleted key: {}, rows affected: {}", key, rowsAffected);
            return rowsAffected > 0;
        }
    }
    
    /**
     * 检查键是否存在
     */
    public boolean exists(String key) throws SQLException {
        String sql = String.format("SELECT 1 FROM %s WHERE key = ?", tableName);
        
        try (PreparedStatement pstmt = connection.prepareStatement(sql)) {
            pstmt.setString(1, key);
            
            try (ResultSet rs = pstmt.executeQuery()) {
                return rs.next();
            }
        }
    }
    
    /**
     * 获取所有键
     */
    public java.util.Set<String> getAllKeys() throws SQLException {
        java.util.Set<String> keys = new java.util.HashSet<>();
        String sql = String.format("SELECT key FROM %s", tableName);
        
        try (Statement stmt = connection.createStatement();
             ResultSet rs = stmt.executeQuery(sql)) {
            
            while (rs.next()) {
                keys.add(rs.getString("key"));
            }
        }
        
        return keys;
    }
    
    /**
     * 清空所有数据
     */
    public void clear() throws SQLException {
        String sql = String.format("DELETE FROM %s", tableName);
        
        try (Statement stmt = connection.createStatement()) {
            int rowsAffected = stmt.executeUpdate(sql);
            logger.info("Cleared {} rows from table {}", rowsAffected, tableName);
        }
    }
    
    /**
     * 获取记录数量
     */
    public long count() throws SQLException {
        String sql = String.format("SELECT COUNT(*) FROM %s", tableName);
        
        try (Statement stmt = connection.createStatement();
             ResultSet rs = stmt.executeQuery(sql)) {
            
            if (rs.next()) {
                return rs.getLong(1);
            }
        }
        
        return 0;
    }
    
    /**
     * 序列化对象
     */
    private byte[] serialize(T obj) throws Exception {
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
             ObjectOutputStream oos = new ObjectOutputStream(baos)) {
            
            oos.writeObject(obj);
            return baos.toByteArray();
        }
    }
    
    /**
     * 反序列化对象
     */
    @SuppressWarnings("unchecked")
    private T deserialize(byte[] data) throws Exception {
        try (ByteArrayInputStream bais = new ByteArrayInputStream(data);
             ObjectInputStream ois = new ObjectInputStream(bais)) {
            
            return (T) ois.readObject();
        }
    }
    
    @Override
    public void close() throws Exception {
        // 不关闭连接，保持连接池
        logger.debug("H2Storage closed (connection kept in pool)");
    }
    
    /**
     * 关闭所有连接池连接
     */
    public static void closeAllConnections() {
        for (Connection conn : connectionPool.values()) {
            try {
                if (conn != null && !conn.isClosed()) {
                    conn.close();
                }
            } catch (SQLException e) {
                logger.error("Failed to close H2 connection", e);
            }
        }
        connectionPool.clear();
        logger.info("All H2 connections closed");
    }
}
